Number of squareful arrays

Time: O(N!); Space: O(N^2); hard

Given an array A of non-negative integers, the array is squareful if for every pair of adjacent elements, their sum is a perfect square.

Return the number of permutations of A that are squareful. Two permutations A1 and A2 differ if and only if there is some index i such that A1[i] != A2[i].

Example 1:

Input: A = [1,17,8]

Output: 2

Explanation:

• [1,8,17] and [17,8,1] are the valid permutations.

Example 2:

Input: [2,2,2]

Output: 1

Notes:

• 1 <= len(A) <= 12

• 0 <= A[i] <= 1e^9

import collections

class Solution1(object):
    """
    Time: O(N!)
    Space: O(N^2)
    """
    def numSquarefulPerms(self, A):
        """
        :type A: List[int]
        :rtype: int
        """
        def dfs(candidate, x, left, count, result):
            count[x] -= 1
            if left == 0:
                result[0] += 1
            for y in candidate[x]:
                if count[y]:
                    dfs(candidate, y, left-1, count, result)
            count[x] += 1

        count = collections.Counter(A)

        candidate = {i: {j for j in count if int((i+j)**0.5) ** 2 == i+j}
                           for i in count}

        result = [0]
        for x in count:
            dfs(candidate, x, len(A)-1, count, result)
        return result[0]

s = Solution1()
A = [1,17,8]
assert s.numSquarefulPerms(A) == 2

See also:

https://leetcode.com/problems/number-of-squareful-arrays